#!/usr/bin/env python
# encoding: utf-8
'''
@author: wangjie
@license: (C) Copyright 2019-2022, Node Supply Chain Manager Corporation Limited.
@contact: wjgd225@sina.com
@software: xuanxun
@file: bbox2mask.py
@time: 2019/12/2 21:24
@desc:
'''
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import numpy as np
import cv2
import json

import xmltodict
try:
    from .cvio import cvio
except:
    import sys
    sys.path.append(os.getcwd())
    from function.cvio import cvio

def readDIR(root):
    ''' Get image list of a root path.

    :param root:
    :return:
    '''
    filelist = []
    for f in os.listdir(root):
        if f.endswith('.jpg') or f.endswith('.jpeg') or f.endswith('.png') or f.endswith('.JPG'):
            #name = os.path.join(root, f)
            filelist.append(f)
    return filelist

def get_mask_dict(img, points, imgpath, jsonname):
    ''' Output the mask label of one image.

    :param img: image np array
    :param points: points list. [[[x1, y1], [x2, y2]...['labelname']]...]
    :param imgpath:
    :param jsonname:
    :return:
    '''
    mask_dict = {}
    mask_dict["imagePath"] = imgpath
    mask_dict["imageHeight"] = img.shape[0]
    mask_dict["imageWidth"]  = img.shape[1]
    mask_dict["imageData"] = None

    mask_dict["lineColor"] = [0, 255, 0, 128]
    mask_dict["fillColor"] = [255, 0, 0, 128]
    mask_dict["flags"] = {}

    shapes = []
    for point in points:
        shape_dict = {}
        shape_dict["line_color"] = None
        shape_dict["fill_color"] = None
        shape_dict["label"] = point[-1]
        shape_dict["shape_type"] = "polygon"
        shape_dict["flags"] = {}
        shape_dict["points"] = point[:-1]

        shapes.append(shape_dict)

    mask_dict["shapes"] = shapes

    with open(jsonname, "w") as f:
        json.dump(mask_dict, f, ensure_ascii=False, indent=4, separators=(',',':'))

def parseXML(xmlpath, bbox2mask_mode):
    ''' parse bbox XML label file.
    :param xmlpath:
    :param bbox2mask_mode: 0: semantic_segmentation, 1:instance_segmentation
    :return:
    '''
    with open(xmlpath, encoding='UTF-8') as fd:
        mn_dict = xmltodict.parse(fd.read())
        
        if isinstance(mn_dict['annotation']['object'], list):
            items = mn_dict['annotation']['object']
        if isinstance(mn_dict['annotation']['object'], dict):
            items = [mn_dict['annotation']['object']]

        labelname = []
        for i, item in enumerate(items):
            name = item['name']
            labelname.append(name)

        labelnameSet = set(labelname)
        labelSetCount = {}
        for labelname in labelnameSet:
            labelSetCount.update({labelname: 0})

        points = []
        for i, item in enumerate(items):
            labelname = item['name']
            xmin = int(float(item['bndbox']['xmin']))
            ymin = int(float(item['bndbox']['ymin']))
            xmax = int(float(item['bndbox']['xmax']))
            ymax = int(float(item['bndbox']['ymax']))
            point1 = [xmin, ymin]
            # point2 = [xmin, int(ymin+(ymax-ymin)/2)]
            point3 = [xmin, ymax]
            # point4 = [int(xmin+(xmax-xmin)/2), ymax]
            point5 = [xmax, ymax]
            # point6 = [xmax, int(ymin+(ymax-ymin)/2)] 
            point7 = [xmax, ymin]
            # point8 = [int(xmin+(xmax-xmin)/2), ymin]
            # 6 points
            # point = [point1, point3, point4, point5, point7, point8]
            # 8 points
            # point = [point1, point2, point3, point4, point5, point6, point7, point8]
            # 4 points
            point = [point1, point3, point5, point7]
            if labelname in labelnameSet:
                if int(bbox2mask_mode) == 0:
                    name = labelname
                elif int(bbox2mask_mode) == 1:
                    labelSetCount[labelname] += 1
                    name = labelname + "-" + str(labelSetCount[labelname])
                else:
                    print("Attention: bbox2mask mode is Not set Right. Current mode is Semantic.")
                    name = labelname
                point.append(name)
            points.append(point)

    return points

import pdb

def bbox2mask(input_path, output_path="", bbox2mask_mode=0):
    ''' To change bbox label to mask label.

    :param input_path: Images and xmls file path.
    :param output_path: Output images and json file path.
    :param bbox2mask_mode: 0: semantic_segmentation, 1:instance_segmentation
    :return:
    '''
    # imglist = readDIR(input_path)
    imglist = cvio.load_image_list(input_path, recursive=False, silent=False)
    n = len(imglist)
    for i, imgfile in enumerate(imglist,1):
        print('[%d/%d][%.2f%%] %s' % (i, n, i / n * 100, os.path.basename(imgfile)))
        name = os.path.splitext(imgfile)[0]
        xmlpath = name + '.xml'
        jsonpath = name + '.json'
        if output_path:
            jsonpath = os.path.join(output_path, os.path.basename(jsonpath))

        image = cv2.imdecode(np.fromfile(imgfile, dtype=np.uint8), cv2.IMREAD_COLOR)
        if os.path.isfile(xmlpath):
            try:
                points = parseXML(xmlpath, bbox2mask_mode)
            except:
                continue
            if len(points) > 0:
                get_mask_dict(image, points, os.path.basename(imgfile), jsonpath)
                # shutil.copy(imgpath, imgpath_out)
            else:
                continue
        else:
            continue

if __name__ == "__main__":
    input_path = r'G:\data\datasets\beer\zjpj\ai\zhujiang_poc_3sku'
    # output_path = os.path.join(input_path, 'mask')
    output_path = r'G:\data\datasets\beer\zjpj\ai\zhujiang_poc_3sku'
    if not os.path.exists(output_path):
        os.mkdir(output_path)
    bbox2mask_mode = 0  # 0: semantic_segmentation, 1:instance_segmentation
    bbox2mask(input_path, output_path, bbox2mask_mode)
